import json
import logging
from channels.generic.websocket import WebsocketConsumer
from asgiref.sync import async_to_sync
from channels.layers import get_channel_layer
import datetime
import time
import threading
import uuid

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    filename='ai_agent.log',
    filemode='a'
)
logger = logging.getLogger('ai_agent')

class AIAgentConsumer(WebsocketConsumer):
    """
    AI代理WebSocket消费者
    提供两个工具：
    1. 分析病例：根据用户提供的身体状况，分析相似病例
    2. 治疗建议：根据病例分析结果，提供是否需要住院和饮食建议
    """
    
    def connect(self):
        """
        建立WebSocket连接时触发
        """
        self.group_name = self.scope['url_route']['kwargs'].get('group', 'default')
        
        # 记录连接日志
        client_info = f"IP: {self.scope['client'][0]}, Port: {self.scope['client'][1]}"
        logger.info(f"有新用户连接: {client_info}, 组名: {self.group_name}")
        print(f"有人进行连接了... 组名: {self.group_name}")
        
        # 将当前连接加入到指定的组
        async_to_sync(self.channel_layer.group_add)(
            str(self.group_name),
            self.channel_name
        )
        
        # 接受WebSocket连接
        self.accept()
        
        # 发送欢迎消息
        self.send(json.dumps({
            'type': 'welcome',
            'message': '欢迎使用AI医疗助手，请描述您的身体状况',
            'session_id': str(uuid.uuid4())
        }))
    
    def disconnect(self, close_code):
        """
        WebSocket连接断开时触发
        """
        logger.info(f"用户断开连接: 组名: {self.group_name}, 关闭代码: {close_code}")
        print(f"连接断开！组名: {self.group_name}")
        
        # 将当前连接从组中移除
        async_to_sync(self.channel_layer.group_discard)(
            str(self.group_name),
            self.channel_name
        )
    
    def receive(self, text_data):
        """
        接收WebSocket消息时触发
        """
        try:
            # 解析接收到的JSON数据
            data = json.loads(text_data)
            message_type = data.get('type', '')
            message_content = data.get('content', '')
            
            logger.info(f"收到消息: {message_type} - {message_content[:50]}...")
            print(f"收到消息: {message_type} - {message_content[:50]}...")
            
            # 根据消息类型处理不同的请求
            if message_type == 'analyze_condition':
                # 工具1: 分析病例
                self._analyze_similar_cases(message_content)
            elif message_type == 'treatment_advice':
                # 工具2: 提供治疗建议
                case_info = data.get('case_info', {})
                self._provide_treatment_advice(case_info)
            else:
                self.send(json.dumps({
                    'type': 'error',
                    'message': '未知的消息类型'
                }))
        except json.JSONDecodeError:
            logger.error(f"无效的JSON数据: {text_data}")
            self.send(json.dumps({
                'type': 'error',
                'message': '无效的JSON数据'
            }))
        except Exception as e:
            logger.error(f"处理消息时出错: {str(e)}")
            self.send(json.dumps({
                'type': 'error',
                'message': f'处理消息时出错: {str(e)}'
            }))
    
    def _analyze_similar_cases(self, condition_description):
        """
        工具1: 分析相似病例
        参数:
            condition_description: 用户描述的身体状况
        """
        # 模拟异步处理，实际项目中应该调用AI大模型API
        def process_analysis():
            # 模拟AI处理延迟
            time.sleep(1)
            
            # 模拟AI分析结果
            # 实际项目中，这里应该调用AI大模型API进行分析
            similar_cases = [
                {
                    'case_id': 'C001',
                    'symptoms': '发热、咳嗽、乏力',
                    'diagnosis': '上呼吸道感染',
                    'similarity': 0.85,
                },
                {
                    'case_id': 'C002',
                    'symptoms': '发热、头痛、肌肉酸痛',
                    'diagnosis': '流行性感冒',
                    'similarity': 0.78,
                }
            ]
            
            # 发送分析结果
            self.send(json.dumps({
                'type': 'analysis_result',
                'message': '根据您的描述，AI分析到以下相似病例',
                'similar_cases': similar_cases,
                'timestamp': datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            }))
        
        # 启动异步处理线程
        thread = threading.Thread(target=process_analysis)
        thread.daemon = True
        thread.start()
        
        # 发送处理中的消息
        self.send(json.dumps({
            'type': 'processing',
            'message': '正在分析您的身体状况，请稍候...'
        }))
    
    def _provide_treatment_advice(self, case_info):
        """
        工具2: 提供治疗建议
        参数:
            case_info: 病例信息
        """
        # 模拟异步处理，实际项目中应该调用AI大模型API
        def process_advice():
            # 模拟AI处理延迟
            time.sleep(1.5)
            
            # 模拟AI分析结果
            # 实际项目中，这里应该调用AI大模型API进行分析
            diagnosis = case_info.get('diagnosis', '')
            symptoms = case_info.get('symptoms', '')
            
            # 根据诊断和症状生成建议
            needs_hospitalization = '发热' in symptoms and ('肺炎' in diagnosis or '感染' in diagnosis)
            
            dietary_advice = [
                "多喝温水，保持水分摄入",
                "清淡饮食，避免辛辣刺激食物",
                "增加蛋白质摄入，如鱼、瘦肉、豆制品等",
                "多吃新鲜蔬果，补充维生素"
            ]
            
            medication_advice = [
                "按医嘱服用药物，不要自行停药或增减剂量",
                "注意观察药物副作用，如有异常及时就医"
            ]
            
            # 发送治疗建议
            self.send(json.dumps({
                'type': 'treatment_advice',
                'needs_hospitalization': needs_hospitalization,
                'hospitalization_reason': "症状较重，需要进一步观察和治疗" if needs_hospitalization else "",
                'dietary_advice': dietary_advice,
                'medication_advice': medication_advice,
                'general_advice': "多休息，保持良好的睡眠，避免过度劳累",
                'timestamp': datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            }))
        
        # 启动异步处理线程
        thread = threading.Thread(target=process_advice)
        thread.daemon = True
        thread.start()
        
        # 发送处理中的消息
        self.send(json.dumps({
            'type': 'processing',
            'message': '正在生成治疗建议，请稍候...'
        }))
    
    # 处理从组内接收到的消息
    def send_message(self, event):
        """
        从组内接收消息并发送给WebSocket客户端
        """
        message = event['message']
        self.send(json.dumps(message)) 